/*
 * anima/x86_64/mm.cc - Memory management functions for x86-64 architecture.
 * Copyright (C) 2007  Mikhail Vorozhtsov
 * See the LICENSE section of the README file for details.
 */

/* $Id$ */

#include HEADER (anima,mm.hh)
#include HEADER (anima,pdb.hh)
#include HEADER (anima,space.hh)
#include HEADER (anima,cpu.hh)
#include HEADER (anima,init.hh)
#include HEADER (anima,console.hh)
#include HEADER (anima,atomic.hh)
#include HEADER (anima/x86_64,pdpm.hh)
#include HEADER (x86_64,regs.hh)
#include HEADER (x86_64,msrs.hh)

START_NAMESPACE2 (anima,x86_64)

static ::x86_64::pe_t kernel_pml4[512] ALIGNED (4096);
static ::x86_64::pe_t kernel_pdp[512] ALIGNED (4096);
static ::x86_64::pe_t kernel_pd[512] ALIGNED (4096);
static ::x86_64::pe_t kernel_pt[512] ALIGNED (4096);

struct allocator_t
{
public:
  INLINE
  void
  install (::x86_64::pe_t pml4e)
  {
    kernel_pml4[0] = pml4e;
  }

  INLINE
  ::x86_64::pe_t *
  map (u64_t paddr)
  {
    mm::map_null (paddr);
    return NULL;
  }

  INLINE
  void
  unmap ()
  {
    mm::unmap_null ();
  }

  INLINE
  ::anima::pdb::lvl1e_t *
  allocate ()
  {
    if (init::done_p ())
      {
        return NULL;
      }
    else
      {
        ::anima::space_t *main_space = init::get_main_space ().get ();

        if (main_space)
          return main_space->allocate (::anima::pdb::entry_t::USAGE_KERNEL);
        else
          return ::anima::pdb::allocate ();
      }
  }

  INLINE
  void
  deallocate (::anima::pdb::lvl1e_t *lvl1e)
  {
    if (init::done_p ())
      {
      }
    else
      {
        ::anima::space_t *main_space = init::get_main_space ().get ();

        if (main_space)
          main_space->deallocate (lvl1e);
        else
          ::anima::pdb::deallocate (lvl1e);
      }
  }
};

static allocator_t kernel_pdpm_allocator;
static self_pdpm_t<GB (511) - MB (2) - KB (4), GB (511) - MB (2), GB (511),
                   allocator_t, true> kernel_pdpm;

extern "C" void *ro_start;
extern "C" void *ro_end;
extern "C" void *bss_end;

bool
mm::init ()
{
  using ::x86_64::pe_t;
  using ::x86_64::vaddr_t;

  // Enable No-execute page-protection.
  ::x86_64::efer_t::get ().set_NXE ().set ();

  // Map kernel.
  {
    u64_t p = 0x100000;

    kernel_pt[((vaddr_t) 0xB8000).PT].set_ptr (0xB8000).set_RW ().set_P ();

    for (; p < (u64_t) &ro_start; p += 4096)
      kernel_pt[((vaddr_t) p).PT].set_ptr (p).set_G ().set_P ();
    for (; p < (u64_t) &ro_end; p += 4096)
      kernel_pt[((vaddr_t) p).PT].set_ptr (p).set_G ().set_NX ().set_P ();
    for (; p < (u64_t) &bss_end; p += 4096)
      kernel_pt[((vaddr_t) p).PT].set_ptr (p)
        .set_G ().set_NX ().set_RW ().set_P ();
  }

  kernel_pd[0].set_ptr ((u64_t) kernel_pt).set_G ().set_RW ().set_P ();
  kernel_pdp[0].set_ptr ((u64_t) kernel_pd).set_G ().set_RW ().set_P ();
  kernel_pml4[0].set_ptr ((u64_t) kernel_pdp).set_G ().set_RW ().set_P ();

  // Save original PML4.
  u64_t orig_pml4 = ::x86_64::cr3_t::get ().PML4;

  // Install new page table.
  ::x86_64::cr3_t::get ().set_ptr ((u64_t) kernel_pml4).set ();

  // Enable write protection.
  ::x86_64::cr0_t::get ().set_WP ().set ();

  // Init kernel PDP mapping.
  if (!kernel_pdpm.init (&kernel_pdpm_allocator, kernel_pml4[0]))
    {
      // Disable write protection.
      ::x86_64::cr0_t::get ().set_WP (false).set ();
      // Restore original page table.
      ::x86_64::cr3_t::get ().set_PML4 (orig_pml4).set ();

      console::error ("Cannot initialize kernel PDP mapping.");
      return false;
    }

  return true;
}

void
mm::map_null (u64_t paddr)
{
  kernel_pt[0].set_raw (0).set_ptr (paddr).set_NX ().set_RW ().set_P ();
}

void
mm::unmap_null ()
{
  kernel_pt[0].set_P (false);
  ::x86_64::inv_tlb_entry (0);
}

END_NAMESPACE2

START_NAMESPACE (anima)

using x86_64::kernel_pml4;
using x86_64::kernel_pdpm;

static atomic64_t last_serial;

bool
mm::init_cpu ()
{
  // Enable No-execute page-protection.
  ::x86_64::efer_t::get ().set_NXE ().set ();

  // Install new page table.
  ::x86_64::cr3_t::get ().set_ptr ((u64_t) kernel_pml4).set ();

  // Enable write protection.
  ::x86_64::cr0_t::get ().set_WP ().set ();

  return true;
}

void
mm::install (space_t *space)
{
  kernel_pml4[1] = space->maps_pml4e;
  kernel_pml4[2] = space->first_pml4e;
  kernel_pml4[3] = space->second_pml4e;

  ::x86_64::cr3_t::get ().set ();
}

void *
mm::map (ct_t ct)
{
  u64_t addr = 0;
  tte_t *pte = kernel_pdpm.map_free (addr);

  if (!pte)
    return NULL;

  pdb::lvl1e_t *lvl1e = NULL;
  
  if (init::done_p ())
    {
    }
  else
    {
      space_t *main_space = init::get_main_space ().get ();

      if (main_space)
        lvl1e = main_space->allocate (pdb::entry_t::USAGE_KERNEL);
      else
        lvl1e = pdb::allocate ();
    }

  if (!lvl1e)
    {
      kernel_pdpm.unmap (addr);
      return NULL;
    }

  pte->wlock ();

  pte->clear ().set_ct (ct).set_base (lvl1e->entry.base)
    .set_NX ().set_RW ().set_P ();

  if (ct != CT_COMMON)
    {
      ((::anima::object_t *) addr)->serial = last_serial.add ();
      ((::anima::object_t *) addr)->paddr = lvl1e->entry.get_ptr ();
    }

  pte->unwlock ();

  return (void *) addr;
}

void *
mm::map (u64_t paddr, tte_t::flag_t flags)
{
  u64_t addr = 0;

  pdb::lvl1e_t *lvl1e = pdb::get (paddr);

  if (lvl1e && lvl1e->entry.usage != pdb::entry_t::USAGE_USER)
    return NULL;

  tte_t *pte = (tte_t *) kernel_pdpm.map_free (addr);

  if (!pte)
    return NULL;

  pte->wlock ();

  pte->clear ().set_ct (CT_COMMON).set_flags (flags)
    .set_ptr (paddr).set_P ();

  pte->unwlock ();

  return (void *) addr;
}

void
mm::unmap (u64_t addr)
{
  ::x86_64::pe_t pte = kernel_pdpm.unmap (addr);

  if (!pte.P)
    return;

  ::x86_64::inv_tlb_entry (addr);

  pdb::lvl1e_t *lvl1e = pdb::get (pte.get_ptr ());

  if (!lvl1e || lvl1e->entry.usage != pdb::entry_t::USAGE_KERNEL)
    return;

  if (init::done_p ())
    {
    }
  else
    {
      space_t *main_space = init::get_main_space ().get ();

      if (main_space)
        main_space->deallocate (lvl1e);
      else
        pdb::deallocate (lvl1e);
    }
}

object_t::state_t
mm::use (object_t *obj, ct_t ct, u64_t serial, object_t::state_t state)
{
  if (state == object_t::STATE_NONE)
    return object_t::STATE_NONE;

  tte_t *pte = kernel_pdpm.get_pte ((u64_t) obj);

  if (!pte)
    return object_t::STATE_NONE;

  rlock (*pte);

  if (!pte->P || pte->get_ct () != ct || !pte->get_I ()
      || (serial != object_t::SERIAL_NONE && obj->serial != serial)
      || obj->refs_cnt.get () == 0)
    {
      pte->unrlock ();
      return object_t::STATE_NONE;
    }

  obj->refs_cnt.add ();

  pte->unrlock ();

  obj->usage_lock.rlock ();

  object_t::state_t result = (object_t::state_t) (obj->state.get () & state);

  if (result == object_t::STATE_NONE)
    {
      obj->usage_lock.unrlock ();
      obj->refs_cnt.sub ();
      return object_t::STATE_NONE;
    }

  return result;
}

void
mm::unuse (object_t *obj)
{
  obj->usage_lock.unrlock ();
  obj->refs_cnt.sub ();
}

object_t::state_t
mm::change_state (object_t *obj, object_t::state_t state)
{
  object_t::state_t result;

  obj->usage_lock.unrlock ();

  obj->usage_lock.wlock ();

  result = obj->state.get ();

  if (result >= state)
    {
      obj->usage_lock.unwlock (result == state);
      return result;
    }

  obj->state.set (state);

  obj->usage_lock.unwlock (true);

  return state;
}

void
mm::activate (object_t *obj)
{
  tte_t *pte = kernel_pdpm.get_pte ((u64_t) obj);

  if (!pte)
    return;

  pte->wlock ();

  if (!pte->P || pte->get_ct () == CT_COMMON || pte->get_I ())
    {
      pte->unwlock ();
      return;
    }

  pte->set_I ();
  obj->state.set (object_t::STATE_ACTIVE);
  obj->usage_lock.rlock ();
  obj->refs_cnt.set (1);

  pte->unwlock ();
}

bool
mm::unmap (object_t *obj, ct_t ct, u64_t serial)
{
  tte_t *pte = kernel_pdpm.get_pte ((u64_t) obj);

  if (!pte)
    return false;

  pte->wlock ();

  if (!pte->P || pte->get_ct () != ct || !pte->get_I ()
      || obj->serial != serial
      || obj->state.get () != object_t::STATE_DEALLOCATING)
    {
      pte->unwlock ();
      return false;
    }

  pte->set_I (false);

  pte->unwlock ();

  while (obj->refs_cnt.get () != 0)
    cpu_t::relax ();

  unmap ((u64_t) obj);

  return true;
}

END_NAMESPACE

